// filepath: ~/stewart_inverse_ws/src/gazebo_joint_control/src/gazebo_joint_control.cpp
#include <gazebo/gazebo.hh>
#include <gazebo/physics/physics.hh>
#include "rclcpp/rclcpp.hpp"
#include <std_msgs/msg/float64_multi_array.hpp>

#include <memory>
#include <vector>

namespace gazebo {

class JointControlPlugin : public ModelPlugin {
public:
  JointControlPlugin() : ModelPlugin() {}

  void Load(physics::ModelPtr _model, sdf::ElementPtr _sdf) override {
    // 检查模型指针是否有效
    if (!_model) {
      RCLCPP_ERROR(rclcpp::get_logger("gazebo_joint_control"),
                   "模型指针无效");
      return;
    }

    this->model_ = _model;

    // 初始化 ROS2 节点
    rclcpp::init(0, nullptr);
    node_ = rclcpp::Node::make_shared("gazebo_joint_control");

    // 获取关节名称列表
    if (_sdf->HasElement("joint_names")) {
      sdf::ElementPtr joint_names_elem = _sdf->GetElement("joint_names");
      sdf::ElementPtr joint_name_elem = joint_names_elem->GetElement("joint_name");
      while (joint_name_elem) {
        std::string joint_name = joint_name_elem->Get<std::string>();
        joint_names_.push_back(joint_name);
        joint_name_elem = joint_name_elem->GetNextElement("joint_name");
      }
    } else {
      RCLCPP_ERROR(rclcpp::get_logger("gazebo_joint_control"),
                   "请在 SDF 文件中指定 <joint_names>");
      return;
    }

    // 获取关节指针
    for (const auto &joint_name : joint_names_) {
      physics::JointPtr joint = model_->GetJoint(joint_name);
      if (!joint) {
        RCLCPP_ERROR(rclcpp::get_logger("gazebo_joint_control"),
                     "找不到关节: %s", joint_name.c_str());
        return;
      }
      joints_.push_back(joint);
    }

    // 订阅 ROS2 主题
    std::string topic_name = "joint_commands";
    if (_sdf->HasElement("topic_name")) {
      topic_name = _sdf->GetElement("topic_name")->Get<std::string>();
    }
    subscription_ = node_->create_subscription<std_msgs::msg::Float64MultiArray>(
        topic_name, 10,
        std::bind(&JointControlPlugin::OnJointCommand, this, std::placeholders::_1));

    // 使用定时器来周期性地执行任务
    update_rate_ = 100.0;
    if (_sdf->HasElement("update_rate")) {
      update_rate_ = _sdf->GetElement("update_rate")->Get<double>();
    }
    update_period_ = std::chrono::duration<double>(1.0 / update_rate_);
    last_update_time_ = node_->now();

    // 连接更新事件
    update_connection_ = gazebo::event::Events::ConnectWorldUpdateBegin(
        std::bind(&JointControlPlugin::OnUpdate, this));

    RCLCPP_INFO(rclcpp::get_logger("gazebo_joint_control"), "插件已加载");
  }

private:
  // ROS2 回调函数，用于接收关节命令
  void OnJointCommand(const std_msgs::msg::Float64MultiArray::SharedPtr msg) {
    if (msg->data.size() != joints_.size()) {
      RCLCPP_WARN(rclcpp::get_logger("gazebo_joint_control"),
                   "接收到的关节命令数量与关节数量不匹配");
      return;
    }

    // 存储最新的关节命令
    latest_joint_commands_ = msg->data;
  }

  // Gazebo 更新回调函数
  void OnUpdate() {
    rclcpp::Time current_time = node_->now();
    rclcpp::Duration time_since_last_update = current_time - last_update_time_;

    if (time_since_last_update >= update_period_) {
      // 设置关节位置
      for (size_t i = 0; i < joints_.size(); ++i) {
        joints_[i]->SetPosition(0, latest_joint_commands_[i]);
      }
      last_update_time_ = current_time;
    }
    // Spin the ROS2 node
    rclcpp::spin_some(node_);
  }

private:
  physics::ModelPtr model_;
  std::vector<physics::JointPtr> joints_;
  std::vector<std::string> joint_names_;

  rclcpp::Node::SharedPtr node_;
  rclcpp::Subscription<std_msgs::msg::Float64MultiArray>::SharedPtr subscription_;

  std::vector<double> latest_joint_commands_;

  gazebo::event::ConnectionPtr update_connection_;

  double update_rate_;
  std::chrono::duration<double> update_period_;
  rclcpp::Time last_update_time_;
};

GZ_REGISTER_MODEL_PLUGIN(JointControlPlugin)

} // namespace gazebo